import cv2
import numpy as np

# 使用的图像
img_path = "./tempImages/3d_img_0.png"

# 矩阵输出到的文件
output_matrix_file_name = "./mat1.txt"

# 绘制当前光标位置和已选的点
def draw_points(img, points, cursor):
    img_copy = img.copy()
    # 画出所有已选的点
    for point in points:
        cv2.circle(img_copy, point, 3, (0, 255, 0), -1)
    # 画出当前光标位置
    cv2.circle(img_copy, cursor, 2, (0, 0, 255), -1)
    coord_text = f"({cursor[0]}, {cursor[1]})"
    cv2.putText(img_copy, coord_text, (cursor[0] + 5, cursor[1] - 5), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.4, (0, 0, 255), 1, cv2.LINE_AA)
    return img_copy

def main():
    # 存储四个点
    points = []

    # points = [
    #     [138, 417],
    #     [502, 417],
    #     [215, 237],
    #     [425, 237]
    # ]

    image = cv2.imread(img_path)
    height, width, channels = image.shape
    print(image.shape[:2])
    clone = image.copy()

    cv2.namedWindow("Select Points", cv2.WINDOW_NORMAL)

    current_point = [int(width/2), int(height/2)]  # 初始光标位置在图像中心

    print("使用 wasd 来移动选点, 用 空格 来确定")

    while True:
        img_with_points = draw_points(clone, points, tuple(current_point))
        cv2.imshow("Select Points", img_with_points)

        key = cv2.waitKey(10) & 0xFF

        # 根据按键调整光标位置
        if key == ord("w"):  # 上
            current_point[1] = max(0, current_point[1] - 1)
        elif key == ord("s"):  # 下
            current_point[1] = min(image.shape[0] - 1, current_point[1] + 1)
        elif key == ord("a"):  # 左
            current_point[0] = max(0, current_point[0] - 1)
        elif key == ord("d"):  # 右
            current_point[0] = min(image.shape[1] - 1, current_point[0] + 1)
        elif key == 32:  # 空格键确定一个点
            if len(points) < 4:
                points.append(tuple(current_point))
                print(f"Point {len(points) - 1} selected at: {current_point}")
            else:
                print("已经选择了4个点。")
        elif key == ord("q"):
            break

        if len(points) == 4:
            break

    print(points)

    # 如果选定了4个点
    if len(points) == 4:
        # 将选定的点转换为浮点数格式
        points = np.float32(points)

        # 重排序
        targets = np.array([[0, 0], [width, 0], [width, height], [0, height]])
        distances = np.linalg.norm(points[:, np.newaxis] - targets, axis=2)
        sorted_indices = np.argmin(distances, axis=0)
        points = points[sorted_indices]
        print("重排序(左上, 右上, 右下, 左下): ")
        print(points)

        h1 = int((points[0][1] + points[1][1]) / 2)
        h2 = int((points[2][1] + points[3][1]) / 2)

        points[0][1] = h1
        points[1][1] = h1
        points[2][1] = h2
        points[3][1] = h2

        # original_middle_point = [int((points[0][0] + points[1][0] + points[2][0] + points[3][0]) / 4), int((h1 + h2) / 2)]

        # for i in range(4):
        #     points[i][0] = points[i][0] - original_middle_point[0]
        #     points[i][1] = points[i][1] - original_middle_point[1]

        pts_dst_size = int(points[2][0] - points[3][0]) # 选用选取的正方形的底边
        # pts_dst_size = int(points[1][0] - points[0][0]) # 选用选取的正方形的顶边


        # 目前采用结果关于y轴(h轴)对称的映射方法
        # 也就是240*320图像中(..., 150)和(..., 170)会被映射为(..., -m)和(..., m)
        pts_dst_origin = {"x": -pts_dst_size/2, "y": -pts_dst_size/2}

        pts_dst = [
            [pts_dst_origin["x"],                   pts_dst_origin["y"]],
            [pts_dst_origin["x"] + pts_dst_size,    pts_dst_origin["y"]],
            [pts_dst_origin["x"] + pts_dst_size,    pts_dst_origin["y"] + pts_dst_size],
            [pts_dst_origin["x"],                   pts_dst_origin["y"] + pts_dst_size]
        ]
        pts_dst = np.float32(pts_dst)

        print("原始点坐标:")
        print(points)
        print("目标点坐标:")
        print(pts_dst)

        M = cv2.getPerspectiveTransform(points, pts_dst)

        print(np.float32(M))
        if isinstance(output_matrix_file_name, str):
            np.savetxt(output_matrix_file_name, np.float32(M), fmt="%.20f")

        # 应用透视变换，将整个图像转换到俯视视角
        warped = cv2.warpPerspective(clone, M, (960, 720))

        # 显示结果
        cv2.imshow("Warped Image", warped)

        while True:
            key = cv2.waitKey(1) & 0xFF
            if key == ord("q"):
                break

    cv2.destroyAllWindows()


if __name__ == "__main__":
    main()